package com.jihu.sharding.algorithm;

import com.google.common.collect.Range;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

/**
 * 数据表复杂分片算法类
 */
public class MyComplexTableShardingAlgorithm implements ComplexKeysShardingAlgorithm<Long> {

    /**
     * 分片
     *
     * @param availableTargetNames 所有可用的真实数据表名称集合
     * @param shardingValue        过滤条件
     * @return 要查询的数据表名集合（支持=或者in）
     */
    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, ComplexKeysShardingValue<Long> shardingValue) {
        // sql: SELECT  cid,cname,user_id,cstatus  FROM course WHERE  cid BETWEEN ? AND ? AND user_id = ?
        // sql: SELECT  cid,cname,user_id,cstatus  FROM course WHERE  cid BETWEEN ? AND ? AND user_id in (?,?)
        // 获取范围字段cid
        Range<Long> cidRange = shardingValue.getColumnNameAndRangeValuesMap().get("cid");
        // 获取精确过滤字段user_id集合
        Collection<Long> userIdCol = shardingValue.getColumnNameAndShardingValuesMap().get("user_id");

        // 获得上限和下限
        Long upperVal = cidRange.upperEndpoint();
        Long lowerVal = cidRange.lowerEndpoint();

        List<String> res = new ArrayList<>();

        for (Long userId : userIdCol) {
            // 实现course_{userID%2+1}逻辑
            BigInteger userIdB = BigInteger.valueOf(userId);
            // 根据算法确定user_id所在的表!!! （再结合数据源的算法, 就可以定位出这行记录所在的位置）
            BigInteger target = (userIdB.mod(new BigInteger("2"))).add(new BigInteger("1"));
            // 返回course_1 / course_2 / course_1,course_2
            res.add(shardingValue.getLogicTableName() + "_" + target);
        }

        return res;
    }
}
